Quench of the singlet fission model starting from a localised singlet state

# pylint: disable=invalid-name
# pylint: disable=too-many-arguments
# pylint: disable=too-many-locals

import os
import os.path

import matplotlib.pyplot as plt
import numpy as np

import qtealeaves as qtl
from qtealeaves.models import get_singlet_fission_1d


def main(
    tn_type=5,
    input_folder=None,
    output_folder=None,
    num_sites=8,
    max_bond_dimension=4,
    timesteps=200,
    plot=True,
):
    """
    Main method for quenching the singlet fission model starting
    from a localised singlet state

    **Arguments**

    tn_type : int, optional
        Choose 5 for python-TTN, 6 for python-MPS.
        Default to 5.

    input_folder : str | None, optional
        Input folder. Default to None.

    output_folder : str | None, optional
        Output folder. Default to None.

    num_sites: int, optional
        number of sites in the system

    max_bond_dimension: int, optional
        maximum bond dimension

    timesteps : int, optional
        Number of timesteps. Default to 160.

    plot : bool, optional
        If True, plot the results at the end of the example. Default tot True
    """
    if input_folder is None:
        input_folder = lambda params: "SF1dDyn/input_L%d" % (params["L"])
    if output_folder is None:
        output_folder = lambda params: "SF1dDyn/output_L%d" % (params["L"])

    model, my_ops = get_singlet_fission_1d()

    my_conv = qtl.convergence_parameters.TNConvergenceParameters(
        max_iter=0,
        max_bond_dimension=max_bond_dimension,
        statics_method=4,
        min_bond_dimension=4,
    )

    my_obs = qtl.observables.TNObservables()
    my_obs += qtl.observables.TNObsLocal("<ns>", "ns")
    my_obs += qtl.observables.TNObsLocal("<nt>", "nt")

    # Dynamics
    quench = qtl.DynamicsQuench("t_grid", time_evolution_mode=2)
    quench["Et"] = lambda tt, params: 0.5
    quench["gamma"] = lambda tt, params: 0.4

    simulation = qtl.QuantumGreenTeaSimulation(
        model,
        my_ops,
        my_conv,
        my_obs,
        tn_type=tn_type,
        folder_name_input=input_folder,
        folder_name_output=output_folder,
        store_checkpoints=False,
    )

    params = []

    # Initial state
    product = [np.array([1.0, 0.0, 0.0], dtype=np.complex128) for _ in range(num_sites)]
    product[0] = np.array([0.0, 1.0, 0.0])
    psi = qtl.emulator.ttn_simulator.TTN.product_state_from_local_states(
        np.array(product),
        convergence_parameters=my_conv,
        tensor_backend=qtl.tensors.TensorBackend(),
    )

    params.append(
        {
            "L": num_sites,
            "Eg": 0.0,
            "Es": 1.0,
            "Et": 0.5,
            "Js": -0.1,
            "Jt": 0.0,
            "chi": 0.0,
            "gamma": 0.0,
            "Quenches": [quench],
            "t_grid": [0.1] * timesteps,
            "continue_file": psi,
            "exclude_from_hash": ["Quenches", "continue_file"],
        }
    )

    simulation.run(params, delete_existing_folder=True)

    for param in params:
        results = simulation.get_dynamic_obs(param)[0]

        tgrid = np.array([np.sum(elem["time"]) for elem in results])
        avg_singlet = np.array([np.sum(elem["<ns>"]) for elem in results])
        avg_triplet = np.array([np.sum(elem["<nt>"]) for elem in results])

        if plot:
            fig = plt.figure()
            ax = fig.add_subplot(111)
            ax.plot(tgrid, avg_singlet, "b-", label="s")
            ax.plot(tgrid, avg_triplet, "r-", label="t")
            ax.set_xlabel("Time (t)")
            ax.set_ylabel("number of singlets, triplets")
            ax.legend()

            out_folder = (
                output_folder(param) if callable(output_folder) else output_folder
            )
            pdf_name = os.path.join(out_folder, "number_over_time.pdf")
            plt.savefig(pdf_name)
            print(f"Saved number of singlets and triplets to {pdf_name}.")

    print("\nExample 'Singletfission_1d_quench.py' ran successfully; ")

    return


if __name__ == "__main__":
    main()